#!/usr/bin/env python3
"""
生成自签名 SSL 证书用于 wss:// 测试
"""

import ssl
from cryptography import x509
from cryptography.x509.oid import NameOID
from cryptography.hazmat.primitives import hashes, serialization
from cryptography.hazmat.primitives.asymmetric import rsa
from datetime import datetime, timedelta, timezone
import os
import ipaddress


def generate_self_signed_cert(hostname="localhost", days=365, extra_ips=None):
    """生成自签名证书"""
    
    # 生成私钥
    private_key = rsa.generate_private_key(
        public_exponent=65537,
        key_size=2048,
    )
    
    # 创建证书
    subject = issuer = x509.Name([
        x509.NameAttribute(NameOID.COUNTRY_NAME, "CN"),
        x509.NameAttribute(NameOID.STATE_OR_PROVINCE_NAME, "Beijing"),
        x509.NameAttribute(NameOID.LOCALITY_NAME, "Beijing"),
        x509.NameAttribute(NameOID.ORGANIZATION_NAME, "Test"),
        x509.NameAttribute(NameOID.COMMON_NAME, hostname),
    ])
    
    # 构建 SAN 列表
    san_list = [
        x509.DNSName(hostname),
        x509.DNSName("127.0.0.1"),
        x509.IPAddress(ipaddress.IPv4Address("127.0.0.1")),
    ]
    
    # 添加额外的 IP 地址
    if extra_ips:
        for ip in extra_ips:
            try:
                san_list.append(x509.IPAddress(ipaddress.IPv4Address(ip)))
                print(f"   添加 IP: {ip}")
            except ValueError:
                print(f"   跳过无效 IP: {ip}")
    
    cert = x509.CertificateBuilder().subject_name(
        subject
    ).issuer_name(
        issuer
    ).public_key(
        private_key.public_key()
    ).serial_number(
        x509.random_serial_number()
    ).not_valid_before(
        datetime.now(timezone.utc)
    ).not_valid_after(
        datetime.now(timezone.utc) + timedelta(days=days)
    ).add_extension(
        x509.SubjectAlternativeName(san_list),
        critical=False,
    ).sign(private_key, hashes.SHA256())
    
    # 保存证书和私钥
    with open("cert.pem", "wb") as f:
        f.write(cert.public_bytes(serialization.Encoding.PEM))
    
    with open("key.pem", "wb") as f:
        f.write(private_key.private_bytes(
            encoding=serialization.Encoding.PEM,
            format=serialization.PrivateFormat.PKCS8,
            encryption_algorithm=serialization.NoEncryption()
        ))
    
    print(f"✅ 已生成自签名证书:")
    print(f"   cert.pem - 证书文件")
    print(f"   key.pem  - 私钥文件")
    print(f"   有效期: {days} 天")
    print(f"   主机名: {hostname}")
    print()
    print("⚠️  注意: 这是自签名证书，浏览器会显示安全警告，测试时请忽略")


if __name__ == "__main__":
    import argparse
    
    parser = argparse.ArgumentParser(description="生成自签名 SSL 证书")
    parser.add_argument("--hostname", default="localhost", help="证书主机名")
    parser.add_argument("--days", type=int, default=365, help="证书有效期（天）")
    parser.add_argument("--extra-ips", nargs="*", help="额外的 IP 地址列表")
    
    args = parser.parse_args()
    
    try:
        generate_self_signed_cert(args.hostname, args.days, args.extra_ips)
    except ImportError:
        print("❌ 缺少依赖: pip install cryptography")
        print("   然后重新运行此脚本")
    except Exception as e:
        print(f"❌ 生成证书失败: {e}")
